import torch
import torch.nn as nn
from mmcls.models.builder import BACKBONES
from mmcls.utils import get_root_logger
from mmcv.runner import load_checkpoint

from .common import ResBlock, conv1x1_block, conv3x3_block, conv7x7_block


class DLAResBlock(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 stride,
                 body_class=ResBlock,
                 return_down=False):
        super(DLAResBlock, self).__init__()
        self.return_down = return_down
        self.downsample = (stride > 1)
        self.project = (in_channels != out_channels)

        self.body = body_class(in_channels=in_channels,
                               out_channels=out_channels,
                               stride=stride)
        self.activ = nn.ReLU(inplace=True)
        if self.downsample:
            self.downsample_pool = nn.MaxPool2d(kernel_size=stride,
                                                stride=stride)
        if self.project:
            self.project_conv = conv1x1_block(in_channels=in_channels,
                                              out_channels=out_channels,
                                              activation=None)

    def forward(self, x):
        down = self.downsample_pool(x) if self.downsample else x
        identity = self.project_conv(down) if self.project else down
        if identity is None:
            identity = x
        x = self.body(x)
        x += identity
        x = self.activ(x)
        if self.return_down:
            return x, down
        else:
            return x


class DLAInitBlock(nn.Module):
    """DLA specific initial block.

    Args:
        in_channels (int):
            Number of input channels.
        out_channels (int):
            Number of output channels
    """
    def __init__(self, in_channels, out_channels):
        super(DLAInitBlock, self).__init__()
        mid_channels = out_channels // 2

        self.conv1 = conv7x7_block(in_channels=in_channels,
                                   out_channels=mid_channels)
        self.conv2 = conv3x3_block(in_channels=mid_channels,
                                   out_channels=mid_channels)
        self.conv3 = conv3x3_block(in_channels=mid_channels,
                                   out_channels=out_channels,
                                   stride=2)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        return x


class DLARoot(nn.Module):
    """
    DLA root block.

    Args:
        in_channels (int):
            Number of input channels.
        out_channels (int):
            Number of output channels.
        residual (bool):
            Whether to use residual connection.
    """
    def __init__(self, in_channels, out_channels, residual):
        super(DLARoot, self).__init__()
        self.residual = residual

        self.conv = conv1x1_block(in_channels=in_channels,
                                  out_channels=out_channels,
                                  activation=None)
        self.activ = nn.ReLU(inplace=True)

    def forward(self, x2, x1, extra):
        last_branch = x2
        x = torch.cat((x2, x1) + tuple(extra), dim=1)
        x = self.conv(x)
        if self.residual:
            x += last_branch
        x = self.activ(x)
        return x


class DLATree(nn.Module):
    def __init__(self,
                 levels,
                 in_channels,
                 out_channels,
                 res_body_class,
                 stride,
                 root_residual,
                 root_dim=0,
                 first_tree=False,
                 input_level=True,
                 return_down=False):
        super(DLATree, self).__init__()
        self.return_down = return_down
        self.add_down = (input_level and not first_tree)
        self.root_level = (levels == 1)

        if root_dim == 0:
            root_dim = 2 * out_channels
        if self.add_down:
            root_dim += in_channels

        if self.root_level:
            self.tree1 = DLAResBlock(in_channels=in_channels,
                                     out_channels=out_channels,
                                     stride=stride,
                                     body_class=res_body_class,
                                     return_down=True)
            self.tree2 = DLAResBlock(in_channels=out_channels,
                                     out_channels=out_channels,
                                     stride=1,
                                     body_class=res_body_class,
                                     return_down=False)
        else:
            self.tree1 = DLATree(levels=levels - 1,
                                 in_channels=in_channels,
                                 out_channels=out_channels,
                                 res_body_class=res_body_class,
                                 stride=stride,
                                 root_residual=root_residual,
                                 root_dim=0,
                                 input_level=False,
                                 return_down=True)
            self.tree2 = DLATree(levels=levels - 1,
                                 in_channels=out_channels,
                                 out_channels=out_channels,
                                 res_body_class=res_body_class,
                                 stride=1,
                                 root_residual=root_residual,
                                 root_dim=root_dim + out_channels,
                                 input_level=False,
                                 return_down=False)
        if self.root_level:
            self.root = DLARoot(in_channels=root_dim,
                                out_channels=out_channels,
                                residual=root_residual)

    def forward(self, x, extra=None):
        extra = [] if extra is None else extra
        x1, down = self.tree1(x)
        if self.add_down:
            extra.append(down)
        if self.root_level:
            x2 = self.tree2(x1)
            x = self.root(x2, x1, extra)
        else:
            extra.append(x1)
            x = self.tree2(x1, extra)
        if self.return_down:
            return x, down
        else:
            return x


@BACKBONES.register_module()
class DLA(nn.Module):
    """DLA model from 'Deep Layer Aggregation'.
       https://arxiv.org/abs/1707.06484.

    Args:
        levels (List[int]):
            Number of levels in each stage.
        channels (List[int]):
            Number of output channels for each stage.
        init_block_channels (int):
            Number of output channels for the initial unit.
        res_body_class (nn.Module):
            Residual block body class.
        residual_root (bool):
            Whether use residual connection in the root blocks.
        in_channels (int, optional):
            Number of input channels. Defaults to 3.
        in_size (tuple, optional):
            Spatial size of the input image. Defaults to (224, 224).
    """
    def __init__(self,
                 levels,
                 channels,
                 init_block_channels,
                 res_body_class=ResBlock,
                 residual_root=False,
                 in_channels=3,
                 in_size=(224, 224)):
        super(DLA, self).__init__()
        self.in_size = in_size

        self.features = nn.Sequential()
        self.features.add_module(
            "init_block",
            DLAInitBlock(in_channels=in_channels,
                         out_channels=init_block_channels))
        in_channels = init_block_channels

        for i in range(len(levels)):
            levels_i = levels[i]
            out_channels = channels[i]
            first_tree = (i == 0)
            self.features.add_module(
                "stage{}".format(i + 1),
                DLATree(levels=levels_i,
                        in_channels=in_channels,
                        out_channels=out_channels,
                        res_body_class=res_body_class,
                        stride=2,
                        root_residual=residual_root,
                        first_tree=first_tree))
            in_channels = out_channels

    def _init_params(self):
        for name, module in self.named_modules():
            if isinstance(module, nn.Conv2d):
                nn.init.kaiming_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)

    def init_weights(self, pretrained):
        if isinstance(pretrained, str):
            logger = get_root_logger()
            load_checkpoint(self, pretrained, strict=False, logger=logger)
        else:
            self._init_params()

    def forward(self, x):
        x = self.features(x)
        return x
